import numpy as np
from calculate_domain_size import nxp1,ny,nz,nx,nyp1,nzp1
from initialize_fdtd_material_grid import eps_r_x,tau_k_x,eps_r_y,tau_k_y,tau_k_z,eps_r_z
from initialize_fdtd_parameters_and_arrays import dt
from initialize_updating_coefficients import psi_x,psi_y,psi_z,den_x,den_y,den_z
# 初始化辅助场量数组
Cqqx = np.zeros((nx, nyp1, nzp1))
Cqqy = np.zeros((nxp1, ny, nzp1))
Cqqz = np.zeros((nxp1, nyp1, nz))

Cqex = np.zeros((nx, nyp1, nzp1))
Cqey = np.zeros((nxp1, ny, nzp1))
Cqez = np.zeros((nxp1, nyp1, nz))

Qxp = np.zeros((nx, nyp1, nzp1))
Qyp = np.zeros((nxp1, ny, nzp1))
Qzp = np.zeros((nxp1, nyp1, nz))

# 计算 x 方向辅助场量更新系数
for ii in range(nx):
    for jj in range(nyp1):
        for kk in range(nzp1):
            if eps_r_x[ii, jj, kk] == 1:
                Cqqx[ii, jj, kk] = 0
                Cqex[ii, jj, kk] = 0
            else:
                Cqqx[ii, jj, kk] = (2 * tau_k_x[ii, jj, kk] - dt) / den_x[ii, jj, kk]
                # Cqqx[ii, jj, kk] = 0  # 注释掉的替代方案
                Cqex[ii, jj, kk] = 2 * psi_x[ii, jj, kk] / den_x[ii, jj, kk]

# 计算 y 方向辅助场量更新系数
for ii in range(nxp1):
    for jj in range(ny):
        for kk in range(nzp1):
            if eps_r_y[ii, jj, kk] == 1:
                Cqqy[ii, jj, kk] = 0
                Cqey[ii, jj, kk] = 0
            else:
                Cqqy[ii, jj, kk] = (2 * tau_k_y[ii, jj, kk] - dt) / den_y[ii, jj, kk]
                # Cqqy[ii, jj, kk] = 0  # 注释掉的替代方案
                Cqey[ii, jj, kk] = 2 * psi_y[ii, jj, kk] / den_y[ii, jj, kk]

# 计算 z 方向辅助场量更新系数
for ii in range(nxp1):
    for jj in range(nyp1):
        for kk in range(nz):
            if eps_r_z[ii, jj, kk] == 1:
                Cqqz[ii, jj, kk] = 0
                Cqez[ii, jj, kk] = 0
            else:
                Cqqz[ii, jj, kk] = (2 * tau_k_z[ii, jj, kk] - dt) / den_z[ii, jj, kk]
                # Cqqz[ii, jj, kk] = 0  # 注释掉的替代方案
                Cqez[ii, jj, kk] = 2 * psi_z[ii, jj, kk] / den_z[ii, jj, kk]
